/*
 * Copyright (c) 2009 Bogong (www.bogong.dk)
 * 
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 * 
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 * 
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
#include <windows.h>

#define LOAD_DRIVER_REG_KEY (_T("Drivers\\LoadDriverTemp"))

static DWORD g_dwOrder = 0xFFFFFFFF;
static DWORD g_dwFlags = 0xFFFFFFFF;
static DWORD g_dwIndex = 0xFFFFFFFF;
static wchar_t g_szPrefix[10] = _T("");
static wchar_t g_szDll[50] = TEXT("");
static wchar_t g_szLoadRegKey[80] = TEXT("");
static wchar_t g_szCustom[10][80];
static bool g_bHelp = false;
static bool g_bList = false;
static DWORD g_dwUnloadId = 0;

typedef enum
{
    TYPE_BOOL,
    TYPE_INT,
    TYPE_DWORD,
    TYPE_STRING
}
OPTION_TYPE;

typedef struct
{
    wchar_t szShort[3];
    wchar_t szLong[20];
    OPTION_TYPE Type;
    void *pVar;
    int uMaxLength;
    bool bAllowMultiple;
    int uCount;
}
CMD_LINE_OPTION;


static CMD_LINE_OPTION ParameterList[] = 
{
    { TEXT("h"),  TEXT("help"),     TYPE_BOOL,   &g_bHelp,         0,                           FALSE, 0},
    { TEXT("i"),  TEXT("list"),     TYPE_BOOL,   &g_bList,         0,                           FALSE, 0},
    { TEXT("f"),  TEXT("flags"),    TYPE_DWORD,  &g_dwFlags,       0,                           FALSE, 0},
    { TEXT("u"),  TEXT("unload"),   TYPE_DWORD,  &g_dwUnloadId,    0,                           FALSE, 0},
    { TEXT("o"),  TEXT("order"),    TYPE_DWORD,  &g_dwOrder,       0,                           FALSE, 0},
    { TEXT("n"),  TEXT("index"),    TYPE_DWORD,  &g_dwIndex,       0,                           FALSE, 0},
    { TEXT("l"),  TEXT("load"),     TYPE_STRING, g_szLoadRegKey,   _countof(g_szLoadRegKey),    FALSE, 0},
    { TEXT("ld"), TEXT("load-dll"), TYPE_STRING, g_szDll,          _countof(g_szDll),           FALSE, 0},
    { TEXT("p"),  TEXT("prefix"),   TYPE_STRING, g_szPrefix,       _countof(g_szPrefix),        FALSE, 0},
    { TEXT("c"),  TEXT("custom"),   TYPE_STRING, g_szCustom,       _countof(g_szCustom[0]),     TRUE,  0},
};
#define ParameterListLength (sizeof(ParameterList) / sizeof(ParameterList[0]))


static void PrintHelp(void)
{
    printf("Options:\n\n");
    printf(" -l <regkey>           Load driver with registry key <regkey>\n");
    printf(" -ld <dll>             Load specified DLL (also requires -p)\n");
    printf(" -u <id>               Unload driver with handle ID <id>\n");
    printf(" -i                    List drivers and their handle IDs\n");
    printf(" -f <flags>            Set flags value (overrides registry if exists)\n");
    printf(" -p <prefix>           Set prefix value (overrides registry if exists)\n");
    printf(" -o <order>            Set order value (overrides registry if exists)\n");
    printf(" -n <index>            Set index value (overrides registry if exists)\n");
    printf(" -c <name>=<value>     Create custom registry key (overrides registry if exists)\n");
    printf("                       Use dword:<num> as value to create a dword. Default is string\n");
    printf(" -h                    This help\n");
    printf("\n");
    printf("Examples:\n");
    printf("\n");
    printf(" To load a driver with settings already specified in registry\n");
    printf(" HKLM\\Drivers\\BuiltIn\\Serial1, use:\n");
    printf("\n");
    printf("    loaddriver -l Serial1\n");
    printf("\n");
    printf(" To load a custom driver with no settings in registry you must\n");
    printf(" specify everything on the command line, e.g.:\n");
    printf("\n");
    printf("    loaddriver -ld serial.dll -p COM -n 4 -c TxBufferSize=dword:1024\n");
    printf("\n");
}




static DWORD AsciiToNumber(wchar_t *pszNumber)
{
    int uRadix = 10;

    // trim leading spaces
    while (*pszNumber == ' ')
    {
        pszNumber++;
    }

    // If it is a hex number, strip leading 0x and set radix to 16 instead of 10
    if ((wcslen(pszNumber) >= 2) && (pszNumber[0] == '0') && ((pszNumber[1] == 'x') || (pszNumber[1] == 'X')))
    {
        uRadix = 16;
        pszNumber += 2;
    }

    // Convert number
    return wcstoul(pszNumber, NULL, uRadix);
}



static bool ParseOption(wchar_t *szOptions[], DWORD &dwIndex, DWORD uMax)
{
    wchar_t *pszName = szOptions[dwIndex];
    int i;
    bool bFound = false;
    bool bRetVal = true;

    if ((pszName) && (*pszName == '-'))
    {
        bool bLongOption = false;
        pszName++; // Skip leading -
        if ((pszName) && (*pszName == '-'))
        {
            bLongOption = true;
            pszName++; // Skip second leading -
        }
        for (i = 0; i < ParameterListLength; i++)
        {
            if (((bLongOption == true) && (wcscmp(pszName, ParameterList[i].szLong) == 0)) ||
                ((bLongOption == false) && (wcscmp(pszName, ParameterList[i].szShort) == 0)))
            {
                bFound = true;
                if ((ParameterList[i].Type != TYPE_BOOL) && (dwIndex + 1 >= uMax))
                {
                    printf("Not enough parameters\n");
                    bRetVal = false;
                    break;
                }
                if (ParameterList[i].bAllowMultiple)
                {
                    ParameterList[i].uCount++;
                }
                else
                {
                    ParameterList[i].uCount = 1;
                }
                switch (ParameterList[i].Type)
                {
                    case TYPE_BOOL:
                        ((bool *)ParameterList[i].pVar)[ParameterList[i].uCount - 1] = true;
                        //printf("--%S=%S\n", ParameterList[i].szLong, *(bool *)ParameterList[i].pVar == true ? TEXT("true") : TEXT("false"));
                        break;
                    case TYPE_INT:
                        ((int *)ParameterList[i].pVar)[ParameterList[i].uCount - 1] = AsciiToNumber(szOptions[++dwIndex]);
                        //printf("--%S=%d\n", ParameterList[i].szLong, *(int *)ParameterList[i].pVar);
                        break;
                    case TYPE_DWORD:
                        ((DWORD *)ParameterList[i].pVar)[ParameterList[i].uCount - 1] = AsciiToNumber(szOptions[++dwIndex]);
                        //printf("--%S=%d\n", ParameterList[i].szLong, *(DWORD *)ParameterList[i].pVar);
                        break;
                    case TYPE_STRING:
                        wcscpy_s(&((wchar_t *)ParameterList[i].pVar)[(ParameterList[i].uCount - 1) * ParameterList[i].uMaxLength], ParameterList[i].uMaxLength, szOptions[++dwIndex]);
                        //printf("--%S=%S, index = %d\n", ParameterList[i].szLong, szOptions[dwIndex], ParameterList[i].uCount);
                        break;
                }
                break;
            }
        }
    }
    if ((bFound == false) && (bRetVal == true))
    {
        printf("Invalid option: %S\n", szOptions[dwIndex]);
        bRetVal = false;
    }
    return bRetVal;
}


static bool ParseOptions(wchar_t *szOptions[], DWORD dwMax)
{
    DWORD i;
    bool bRetVal = true;

    memset(g_szCustom, 0, sizeof(g_szCustom));
    for (i = 1; i < dwMax; i++)
    {
        bRetVal = ParseOption(szOptions, i, dwMax);
        if (bRetVal != true)
        {
            break;
        }
    }
    return bRetVal;
}


static bool LoadDriver(wchar_t *pszRegKey)
{
    bool bRetVal = true;
    HANDLE hDriver = ActivateDeviceEx(pszRegKey, NULL, 0, NULL);
    if ((hDriver != INVALID_HANDLE_VALUE) && (hDriver != 0))
    {
        printf("Driver loaded. Handle id = 0x%08X\n", (DWORD)hDriver);
        CloseHandle(hDriver);
    }
    else
    {
        // Note: the driver may have succeded to load if the handle is zero
        printf("Failed to activate driver: %d\n", GetLastError());
        bRetVal = false;
    }
    return bRetVal;
}

static void UnloadDriver(HANDLE hDriver)
{
    printf("Unloading driver with handle 0x%X\n", (DWORD)hDriver);
    if (DeactivateDevice(hDriver) == TRUE)
    {
        printf("Driver unloaded\n");
    }
    else
    {
        printf("Failed to unload driver: %d\n", GetLastError());
    }
}

static void ListRunningDrivers( void )
{
    HANDLE hDriver;
    DeviceSearchType DeviceSearch = DeviceSearchByLegacyName;
    DEVMGR_DEVICE_INFORMATION DeviceInfo;

    DeviceInfo.dwSize = sizeof( DEVMGR_DEVICE_INFORMATION );
    hDriver = FindFirstDevice(DeviceSearch, L"*", &DeviceInfo);

    if( hDriver != INVALID_HANDLE_VALUE )
    {
        printf("%-6s %-10s %-12s %-45s %-30s\n", "Legacy", "Handle", "Name", "Key", "BUS");
        do
        {
            printf("%-6S 0x%08X %-12S %-45S %-30S\n",
                DeviceInfo.szLegacyName,
                (DWORD)DeviceInfo.hDevice,
                DeviceInfo.szDeviceName,
                DeviceInfo.szDeviceKey,
                DeviceInfo.szBusName);

        } 
        while( FindNextDevice( hDriver, &DeviceInfo ) );
        FindClose( hDriver );
    }
    else
    {
        printf("Error opening device list: %d\n", GetLastError());
    }
}


static bool SetRegKeys(wchar_t *pszRegKey, wchar_t *pszDll, wchar_t *pszPrefix, DWORD dwIndex, DWORD dwOrder, DWORD dwFlags)
{
    HKEY hTargetKey;
    DWORD dwDisposition;
    DWORD RetVal;
    DWORD uSize;
    bool bRetVal = true;

    // Create key if it does not already exist
    RetVal = RegCreateKeyEx(HKEY_LOCAL_MACHINE, pszRegKey, 0, NULL, REG_OPTION_VOLATILE, KEY_ALL_ACCESS, NULL, &hTargetKey, &dwDisposition);

    // If it opened, then create keys
    if (RetVal == ERROR_SUCCESS)
    {
        if (dwOrder != 0xFFFFFFFF)
        {
            uSize = sizeof(DWORD);
            RegSetValueEx(hTargetKey, _T("Order"), 0, REG_DWORD, (CONST BYTE *)&dwOrder, uSize);
            printf(" - Order = %d\n", dwOrder);
        }
        if (dwIndex != 0xFFFFFFFF)
        {
            uSize = sizeof(DWORD);
            RegSetValueEx(hTargetKey, _T("Index"), 0, REG_DWORD, (CONST BYTE *)&dwIndex, uSize);
            printf(" - Index = %d\n", dwIndex);
        }
        if (dwFlags != 0xFFFFFFFF)
        {
            uSize = sizeof(DWORD);
            RegSetValueEx(hTargetKey, _T("Flags"), 0, REG_DWORD, (CONST BYTE *)&dwFlags, uSize);
            printf(" - Flags = 0x%X\n", dwFlags);
        }
        if (wcslen(pszDll) > 0)
        {
            uSize = (wcslen(pszDll) + 1) * sizeof(wchar_t);
            RegSetValueEx(hTargetKey, _T("Dll"), 0, REG_SZ, (CONST BYTE *)pszDll, uSize);
            printf(" - Dll = %S\n", pszDll);
        }
        if (wcslen(pszPrefix) > 0)
        {
            uSize = (wcslen(pszPrefix) + 1) * sizeof(wchar_t);
            RegSetValueEx(hTargetKey, _T("Prefix"), 0, REG_SZ, (CONST BYTE *)pszPrefix, uSize);
            printf(" - Prefix = %S\n", pszPrefix);
        }
        RegCloseKey(hTargetKey);
    }
    else
    {
        printf("Failed to open registry key: %d\n", GetLastError());
        bRetVal = false;
    }
    return bRetVal;
}


static bool SetCustomRegKey(wchar_t *pszRegKey, wchar_t *pszKeyValue)
{
    wchar_t *pszContext;
    wchar_t *pszKey = wcstok_s(pszKeyValue, TEXT("="), &pszContext);
    bool bRetVal = false;
    if (pszKey != NULL)
    {
        wchar_t *pszValue = wcstok_s(NULL, TEXT("="), &pszContext);
        if (pszValue != NULL)
        {
            HKEY hTargetKey;
            DWORD dwDisposition;
            DWORD RetVal;
            DWORD uSize;

            bRetVal = true;

            // Create key if it does not already exist
            RetVal = RegCreateKeyEx(HKEY_LOCAL_MACHINE, pszRegKey, 0, NULL, REG_OPTION_VOLATILE, KEY_ALL_ACCESS, NULL, &hTargetKey, &dwDisposition);

            // If it opened, then create keys
            if (RetVal == ERROR_SUCCESS)
            {
                if (wcsncmp(pszValue, TEXT("dword:"), 6) == 0)
                {
                    DWORD dwValue = AsciiToNumber(&pszValue[6]);
                    uSize = sizeof(DWORD);
                    RegSetValueEx(hTargetKey, pszKey, 0, REG_DWORD, (CONST BYTE *)&dwValue, uSize);
                    printf(" - %S = %d\n", pszKey, dwValue);
                }
                else
                {
                    uSize = (wcslen(pszValue) + 1) * sizeof(wchar_t);
                    RegSetValueEx(hTargetKey, pszKey, 0, REG_SZ, (CONST BYTE *)pszValue, uSize);
                    printf(" - %S = %S\n", pszKey, pszValue);
                }
                RegCloseKey(hTargetKey);
            }
        }
    }
    return bRetVal;
}


static bool Execute()
{
    bool bRetVal = true;

    if (g_bHelp == true)
    {
        PrintHelp();
    }
    else if (g_bList == true)
    {
        ListRunningDrivers();
    }
    else
    {
        wchar_t szRegKey[80] = TEXT("");
        bool bLoadDriver = false;

        if (wcslen(g_szLoadRegKey) > 0)
        {
            if (wcschr(g_szLoadRegKey, '\\') == NULL)
            {
                wcscat_s(szRegKey, _countof(szRegKey), TEXT("Drivers\\BuiltIn\\"));
            }
            wcscat_s(szRegKey, _countof(szRegKey), g_szLoadRegKey);
            bLoadDriver = true;
            printf("Loading driver with settings from registry: %S\n", szRegKey);
        }
        else if ((wcslen(g_szDll) > 0) && (wcslen(g_szPrefix) > 0))
        {
            wcscat_s(szRegKey, _countof(szRegKey), LOAD_DRIVER_REG_KEY);
            bLoadDriver = true;
            printf("Loading driver with custom settings\n");
        }
        else if ((wcslen(g_szDll) > 0) || (wcslen(g_szPrefix) > 0))
        {
            printf("You have to specify both the DLL and the prefix to load a custom DLL\n");
            bRetVal = false;
        }

        if (bRetVal == true)
        {
            int i;

            // Set reg keys
            SetRegKeys(szRegKey, g_szDll, g_szPrefix, g_dwIndex, g_dwOrder, g_dwFlags);

            // Set custom reg keys
            for (i = 0; i < _countof(g_szCustom); i++)
            {
                if (wcslen(g_szCustom[i]) > 0)
                {
                    SetCustomRegKey(szRegKey, g_szCustom[i]);
                }
            }

            if (bLoadDriver == true)
            {
                LoadDriver(szRegKey);
            }
            else if (g_dwUnloadId > 0)
            {
                UnloadDriver((HANDLE)g_dwUnloadId);
            }
        }
    }
    return bRetVal;
}


int wmain(int argc, wchar_t* argv[])
{
    bool bRetVal = false;

    printf("Windows CE LoadDriver 0.2 (Compile time: %s %s)\n", __DATE__, __TIME__);
    printf("Copyright (C) 2009 by Bogong (www.bogong.dk)\n\n");
    if (argc > 1)
    {
        bRetVal = ParseOptions(argv, argc);
        if (bRetVal == true)
        {
            bRetVal = Execute();
        }
    }
    else
    {
        PrintHelp();
    }
    return (bRetVal == true ? 0 : 1);
}
